import torch
from tqdm import tqdm
from utils.config import *

# 验证
def validate(model, validate_loader, epoch, writer):
    num_batches = len(validate_loader)
    avg_error_avg = 0.0
    avg_error_1px = 0.0
    avg_error_2px = 0.0
    avg_error_3px = 0.0
    avg_error_4px = 0.0
    avg_error_5px = 0.0

    for  batch in tqdm(validate_loader): # 创建进度条
        # 从batch中取样本
        left_img = batch['left'].to(device)
        right_img = batch['right'].to(device)
        target_disp = batch['disp'].to(device)

        # 掩码
        mask = (target_disp < max_disp) & (target_disp > 0)
        mask = mask.detach_()

        # 前向传播
        with torch.no_grad():
            disp = model(left_img, right_img)

        # 计算batch的误差
        epe = torch.abs(disp[mask] - target_disp[mask]) # 预测视差和真实视差的差值
        # error_mask = (epe >= 3.0) & (epe >= (target_disp[mask] * 0.05)) # 掩码，误差绝对值>3px且相对误差>5%
        error_avg = torch.sum(epe).item() / torch.numel(disp[mask])
        error_1px = torch.sum((epe >= 1.0)).item() / torch.numel(disp[mask])
        error_2px = torch.sum((epe >= 2.0)).item() / torch.numel(disp[mask])
        error_3px = torch.sum((epe >= 3.0)).item() / torch.numel(disp[mask])
        error_4px = torch.sum((epe >= 4.0)).item() / torch.numel(disp[mask])
        error_5px = torch.sum((epe >= 5.0)).item() / torch.numel(disp[mask])

        avg_error_avg += error_avg
        avg_error_1px += error_1px
        avg_error_2px += error_2px
        avg_error_3px += error_3px
        avg_error_4px += error_4px
        avg_error_5px += error_5px

    avg_error_avg = avg_error_avg / num_batches
    avg_error_1px = avg_error_1px / num_batches * 100
    avg_error_2px = avg_error_2px / num_batches * 100
    avg_error_3px = avg_error_3px / num_batches * 100
    avg_error_4px = avg_error_4px / num_batches * 100
    avg_error_5px = avg_error_5px / num_batches * 100

    print('epoch: {:03} | avg_error: {:5.3} | 1px-error: {:5.3}% | 2px-error: {:5.3}% | 3px-error: {:5.3}% | 4px-error: {:5.3}% | 5px-error: {:5.3}%'.format(epoch, avg_error_avg, avg_error_1px, avg_error_2px, avg_error_3px, avg_error_4px, avg_error_5px))

    # 绘制avg_error图像
    writer.add_scalar("avg_error_avg", avg_error_avg, epoch)
    writer.add_scalar("avg_error_1px", avg_error_1px, epoch)
    writer.add_scalar("avg_error_2px", avg_error_2px, epoch)
    writer.add_scalar("avg_error_3px", avg_error_3px, epoch)
    writer.add_scalar("avg_error_4px", avg_error_4px, epoch)
    writer.add_scalar("avg_error_5px", avg_error_5px, epoch)

    return {'avg_error': avg_error_avg,
            'error_1px': avg_error_1px,
            'error_2px': avg_error_2px,
            'error_3px': avg_error_3px,
            'error_4px': avg_error_4px,
            'error_5px': avg_error_5px}
